"""Module for time step functions."""

import numpy

class BurgersTimeStep(object):
    def __init__(self, grid):
        self.grid = grid

    def get_max_speed(self):
        return numpy.max( self.grid.q[0] )

class IETimeStep(object):
    def __init__(self, a):
        self.a = a

    def get_max_speed(self):
        return self.a

class EulerEquationsTimeStep(object):
    def __init__(self, grid):
        self.grid = grid

    def get_max_speed(self):
        """Compute the time step.

        The formula used is eqn (6.13) and (6.20) from Toro's book:

        ..math::

        \Delta t \leq \frac{\Delta x}{S^n_{max}},

        S^n_{max} = max{ |u_i| + a_i }
        
        """
        grid = self.grid

        # get the particle velocities and sound speed
        u = grid.w[1]; a = grid.a

        smax = max( a + u )
        return smax


